﻿using System;
using System.Collections.Generic;

namespace DataStructures.Heap
{
    /// <summary>
    ///     A generic implementation of a binary heap.
    /// </summary>
    /// <remarks>
    ///     A binary heap is a complete binary tree that satisfies the heap property;
    ///     that is every node in the tree compares greater/less than or equal to its left and right
    ///     child nodes. Note that this is different from a binary search tree, where every node
    ///     must be the largest/smallest node of all of its children.
    ///     Although binary heaps are not very efficient, they are (probably) the simpliest heaps
    ///     to understand and implement.
    ///     More information: https://en.wikipedia.org/wiki/Binary_heap .
    /// </remarks>
    /// <typeparam name="T">Type of elements in binary heap.</typeparam>
    public class BinaryHeap<T>
    {
        /// <summary>
        ///     Comparer to use when comparing elements.
        /// </summary>
        private readonly Comparer<T> comparer;

        /// <summary>
        ///     List to hold the elements of the heap.
        /// </summary>
        private readonly List<T> data;

        /// <summary>
        ///     Initializes a new instance of the <see cref="BinaryHeap{T}" /> class.
        /// </summary>
        public BinaryHeap()
        {
            data = new List<T>();
            comparer = Comparer<T>.Default;
        }

        /// <summary>
        ///     Initializes a new instance of the <see cref="BinaryHeap{T}" /> class with a custom comparision function.
        /// </summary>
        /// <param name="customComparer">The custom comparing function to use to compare elements.</param>
        public BinaryHeap(Comparer<T> customComparer)
        {
            data = new List<T>();
            comparer = customComparer;
        }

        /// <summary>
        ///     Gets the number of elements in the heap.
        /// </summary>
        public int Count => data.Count;

        /// <summary>
        ///     Add an element to the binary heap.
        /// </summary>
        /// <remarks>
        ///     Adding to the heap is done by append the element to the end of the backing list,
        ///     and pushing the added element up until the heap property is restored.
        /// </remarks>
        /// <param name="element">The element to add to the heap.</param>
        /// <exception cref="ArgumentException">Thrown if element is already in heap.</exception>
        public void Push(T element)
        {
            data.Add(element);
            HeapifyUp(data.Count - 1);
        }

        /// <summary>
        ///     Remove the top/root of the binary heap (ie: the largest/smallest element).
        /// </summary>
        /// <remarks>
        ///     Removing from the heap is done by swapping the top/root with the last element in
        ///     the backing list, removing the last element, and pushing the new root down
        ///     until the heap property is restored.
        /// </remarks>
        /// <returns>The top/root of the heap.</returns>
        /// <exception cref="InvalidOperationException">Thrown if heap is empty.</exception>
        public T Pop()
        {
            if (Count == 0)
            {
                throw new InvalidOperationException("Heap is empty!");
            }

            var elem = data[0];
            data[0] = data[^1];
            data.RemoveAt(data.Count - 1);
            HeapifyDown(0);

            return elem;
        }

        /// <summary>
        ///     Return the top/root of the heap without removing it.
        /// </summary>
        /// <returns>The top/root of the heap.</returns>
        /// <exception cref="InvalidOperationException">Thrown if heap is empty.</exception>
        public T Peek()
        {
            if (Count == 0)
            {
                throw new InvalidOperationException("Heap is empty!");
            }

            return data[0];
        }

        /// <summary>
        ///     Returns element if it compares larger to the top/root of the heap, else
        ///     inserts element into the heap and returns the top/root of the heap.
        /// </summary>
        /// <param name="element">The element to check/insert.</param>
        /// <returns>element if element compares larger than top/root of heap, top/root of heap otherwise.</returns>
        public T PushPop(T element)
        {
            if (Count == 0)
            {
                return element;
            }

            if (comparer.Compare(element, data[0]) < 0)
            {
                var tmp = data[0];
                data[0] = element;
                HeapifyDown(0);
                return tmp;
            }

            return element;
        }

        /// <summary>
        ///     Check if element is in the heap.
        /// </summary>
        /// <param name="element">The element to check for.</param>
        /// <returns>true if element is in the heap, false otherwise.</returns>
        public bool Contains(T element) => data.Contains(element);

        /// <summary>
        ///     Remove an element from the heap.
        /// </summary>
        /// <remarks>
        ///     In removing an element from anywhere in the heap, we only need to push down or up
        ///     the replacement value depending on how the removed value compares to its
        ///     replacement value.
        /// </remarks>
        /// <param name="element">The element to remove from the heap.</param>
        /// <exception cref="ArgumentException">Thrown if element is not in heap.</exception>
        public void Remove(T element)
        {
            var idx = data.IndexOf(element);

            if (idx == -1)
            {
                throw new ArgumentException($"{element} not in heap!");
            }

            Swap(idx, data.Count - 1);
            var tmp = data[^1];
            data.RemoveAt(data.Count - 1);

            if (idx < data.Count)
            {
                if (comparer.Compare(tmp, data[idx]) > 0)
                {
                    HeapifyDown(idx);
                }
                else
                {
                    HeapifyUp(idx);
                }
            }
        }

        /// <summary>
        ///     Swap the elements in the heap array at the given indices.
        /// </summary>
        /// <param name="idx1">First index.</param>
        /// <param name="idx2">Second index.</param>
        private void Swap(int idx1, int idx2)
        {
            var tmp = data[idx1];
            data[idx1] = data[idx2];
            data[idx2] = tmp;
        }

        /// <summary>
        ///     Recursive function to restore heap properties.
        /// </summary>
        /// <remarks>
        ///     Restores heap property by swapping the element at <paramref name="elemIdx" />
        ///     with its parent if the element compares greater than its parent.
        /// </remarks>
        /// <param name="elemIdx">The element to check with its parent.</param>
        private void HeapifyUp(int elemIdx)
        {
            var parent = (elemIdx - 1) / 2;

            if (parent >= 0 && comparer.Compare(data[elemIdx], data[parent]) > 0)
            {
                Swap(elemIdx, parent);
                HeapifyUp(parent);
            }
        }

        /// <summary>
        ///     Recursive function to restore heap properties.
        /// </summary>
        /// <remarks>
        ///     Restores heap property by swapping the element at <paramref name="elemIdx" />
        ///     with the larger of its children.
        /// </remarks>
        /// <param name="elemIdx">The element to check with its children.</param>
        private void HeapifyDown(int elemIdx)
        {
            var left = 2 * elemIdx + 1;
            var right = 2 * elemIdx + 2;

            var leftLargerThanElem = left < Count && comparer.Compare(data[elemIdx], data[left]) < 0;
            var rightLargerThanElem = right < Count && comparer.Compare(data[elemIdx], data[right]) < 0;
            var leftLargerThanRight = left < Count && right < Count && comparer.Compare(data[left], data[right]) > 0;

            if (leftLargerThanElem && leftLargerThanRight)
            {
                Swap(elemIdx, left);
                HeapifyDown(left);
            }
            else if (rightLargerThanElem && !leftLargerThanRight)
            {
                Swap(elemIdx, right);
                HeapifyDown(right);
            }
        }
    }
}
